###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*##
###*### FUNCTION THAT SIMULATES A MARKET EQUILIBRIUM        ###*###*###*###*###*###*###*##
###*### FOR THE MODEL OF CHILD CARE AND LABOR FORCE CHOICE  ###*###*###*###*###*###*###*##
###*### GIVEN A PARAMETER VECTOR                            ###*###*###*###*###*###*###*##
###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*###*##

simul_market <- function(parameters, both = FALSE, mini = 0, tau = 0, 
                         Cash = 0, hqv = FALSE, hqr = FALSE, local = F) {
  # _________________________________________________________________________________ ####
  # PRIMITIVE PARAMETERS TO ESTIMATE ####
  #' Assigns the values of the input vector to the model parameters
  
  ## Apply constraints
  pars <- conv(parameters)
  pars <- constraints(pars)
  pars <- conv.inv(pars, local = local)
  
  ## Household utility
  ma1  <- pars["ma1"]
  ma2  <- pars["ma2"]
  ma3  <- pars["ma3"]
  sa1  <- pars["sa1"]
  sa2  <- pars["sa2"]
  sa3  <- pars["sa3"]
  ma1s <- pars["ma1s"]
  ma2s <- pars["ma2s"]
  sa1s <- pars["sa1s"]
  sa2s <- pars["sa2s"]
  mu   <- pars["mu"]
  
  ## Child quality
  g0  <- pars["g0"]
  g0s <- pars["g0s"]
  gf  <- pars["gf"]
  gm  <- pars["gm"]
  gmh <- pars["gmh"]
  gmc <- pars["gmc"]
  gsm <- pars["gsm"]
  gR  <- pars["gR"] 
  gN  <- pars["gN"] 
  gL  <- pars["gL"] 
  gH  <- pars["gH"] 
  re  <- pars["re"]
  
  ## Wages offer
  rho  <- pars["rho"]
  rhoc <- pars["rhoc"]
  mwm  <- pars["mwm"]
  swm  <- pars["swm"]
  mwmh <- pars["mwmh"] 
  swmh <- pars["swmh"]
  mwmc <- pars["mwmc"] 
  swmc <- pars["swmc"]
  
  ## Childcare costs
  mcL <- pars["mcL"]
  mcH <- pars["mcH"]
  scL <- pars["scL"]
  scH <- pars["scH"]
  tcc <- pars["tcc"]
  
  # _________________________________________________________________________________ ####
  # HETEROGENEITY SOURCES ####
  #' Defines the unobserved heterogeneity sources and wage offers
  
  ## Utility heterogeneity
  v  <- cbind(ev[, 1]*sa1 + ma1,
              ev[, 2]*sa2 + ma2, 
              ev[, 3]*sa3 + ma3)
  vs <- cbind(ev[, 1]*sa1s + ma1s,
              ev[, 2]*sa2s + ma2s, 
              ev[, 3]*sa3 + ma3)
  v  <- v*(1 - sm) + vs*sm
  # Single mom fix
  ai <- cbind(exp(v[1:J, ]), 1)
  ai[sm == 1, 3] <- 0
  ai <- ai/rowSums(ai)
  
  ## Production constant
  g0i <- g0 + g0s*sm
  
  ## Wage offers
  mwmi <- mwm*(ed == 1) + mwmh*(ed == 2) + mwmc*(ed == 3)
  swmi <- swm*(ed == 1) + swmh*(ed == 2) + swmc*(ed == 3)
  rhoi <- rho*(ed <= 2) + rhoc*(ed == 3)
  wi            <- exp(mwmi + swmi*ew + rhoi*ef)
  wi[wi < minw] <- minw
  wi            <- cbind(wf, wi)
  
  ## Childcare costs
  Cg <- exp(cbind(mcL + ec*scL, mcL + ec*scH))
  
  # _________________________________________________________________________________ ####
  # HOUSEHOLD OPTIONS ####
  #' Defining households' choice set, endowments and constraints
  #' as well as computing the dad's optimal time allocation for every labor and
  #' child care arrangement
  
  ## Household options
  hhopts <- data.frame(ha = 8, hb = c(8, 8, 8, 8, 
                                      4, 4, 4, 4, 4,
                                      0, 0, 0, 0, 0), 
                       gd = c(gL, gH, gN, gR, 
                              gL, gH, gN, gR, 0, 
                              gL, gH, gN, gR, 0))
  
  # Household options
  Xa  <- cbind(q0, ai[, 2], ai[, 3], ai[, 4], ed)[sm == 0, ]
  X   <- suppressWarnings(t(apply(Xa, 1, function(x) {
    # Define tstar for each scenario
    out <- apply(hhopts, 1, function(xk) {
      # print(x)
      # print(xk)
      gmi  <- gm*(x[5] == 1) + gmh*(x[5] == 2) + gmc*(x[5] == 3)
      anst <- optim(par = 1, fn = uta, method = "Nelder-Mead",
                    ha = xk[1], hb = xk[2], gd = xk[3], q0 = x[1], 
                    a2 = x[2], a3 = x[3], a4 = x[4], g0 = g0, gf = gf, 
                    gm = gmi, re = re, gL = gL, gH = gH, gN = gN, tcc = tcc)$par
      cc   <- xk[2] == 4 & xk[3] == 0
      cd   <- xk[2] == 8 & xk[3] != gN  #xk[3] == gL | xk[3] == gH
      anst <- max(0.01 + 4*cc + tcc*cd, min(7.99, anst))
      return(anst)
    })
    return(out)
  })))
  tstar <- matrix(0, J, ncol(X))
  tstar[sm == 0, ] <- X
  colnames(tstar)  <- c("8,8,L", "8,8,H", "8,8,N", "8,8,R", 
                        "8,4,L", "8,4,H", "8,4,N", "8,4,R", "8,4,0",
                        "8,0,L", "8,0,H", "8,0,N", "8,0,R", "8,0,0")
  
  # Time allocations
  hb <- cbind(matrix(c(8, 4, 0), J, 3*(NL + 2), byrow = TRUE), 4, 0)
  ha <- hb*0 + 8*(1 - sm)
  
  # BL income, subs, and discounts
  inco   <- wi[, 1]*ha + wi[, 2]*hb + I*sm + Cash
  smdisc <- subs*(sm*(ed <= 2) + hb*0)
  
  # CF discounts: center vouchers
  dmini  <- if(mini > 0) 1*(inco <= mini) else 1
  ht     <- if(both) 1*(hb > 0) else 1 + hb*0
  c_disc <- tau*ht*dmini
  c_disc[, (3*NL):ncol(c_disc)] <- 0
  
  # _________________________________________________________________________________ ####
  # COMPUTE THE SP-NASH EQUILIBRIUM ####
  #' Computes the investment (1st) and respective competition (2nd) stages 
  #' of the supply side 2-stage game.
  
  ## Compute equilibrium
  EntOrder <- order(Cg[, 1])
  EQ2 <- sfLapply(1:(3^NL), function(h) {
    #EQ2 <- list()
    #for(h in 1:(3^NL))  {
    #  print(h)
    g <- gstate(h, gL, gH)
    s <- g == 0
    if(h == 1 | sum(seq_len(sum(s)) == which(s)) == sum(s) & sum(s) < NL) {
      g   <- g[NL:1]
      g   <- g[order(EntOrder)]
      out <- master(g, Cg, Cn, Cr, Fixed, gL, gH, gN, gR, ha, hb, 
                    q0, g0i, gf, gm, gmh, gmc, gsm, re, tstar, hhopts, 
                    ed, inco, c_disc, smdisc, hqv, reli, 
                    ai[, 1], ai[, 2], ai[, 3], ai[, 4], mu, tcc)$out
    } else {
      g   <- g[NL:1]
      g   <- g[order(EntOrder)]
      out <- data.frame(g = c(g, gN, gR), p = 0, pi = -1, s = 0, mc = 0)
    }
    #EQ2[[h]] <- out
    #}
    return(out)
  })
  EQ  <- EQ2
  if(hqr) {
    EQ <- lapply(EQ, function(x) {
      x$pi[x$g == gL] <- -1000
      return(x)
    })
  }
  
  ## Compute the SP equilibrium with backward induction
  Y <- sapply(EQ, function(x) x$pi[1:NL])
  #G <- sapply(EQ, function(x) x$g[1:NL])
  colnames(Y) <- 1:length(EQ)
  for(n in EntOrder[NL:2]) {
    if(n == EntOrder[NL]) y <- Y
    x <- apply(matrix(y[n, ], ncol = 3, byrow = TRUE), 1, which.max)
    x <- 3*(1:length(x)) + x - 3
    y <- y[, x]
  }
  # Define leader's decision
  ind  <- which.max(y[EntOrder[1], ])
  ind  <- as.numeric(colnames(y)[ind])
  SPEQ <- EQ[[ind]]
  
  # _________________________________________________________________________________ ####
  # EQUILIBRIUM VARIABLES ####
  #' Define the relevant outcome variables from the computed equilibrium to compute 
  #' the simulated moments
  
  ## Equilibrium variables
  eq <- master(SPEQ$g[1:NL], Cg, Cn, Cr, Fixed, gL, gH, gN, gR, ha, hb, 
               q0, g0i, gf, gm, gmh, gmc, gsm, re, tstar, hhopts, ed, inco, 
               c_disc, smdisc, hqv, reli, ai[, 1], ai[, 2], ai[, 3], ai[, 4], mu, tcc)
  SPEQ <- eq$out
  p    <- SPEQ$p
  g    <- SPEQ$g
  gmat <- matrix(c(rep(g, each = 3), 0, 0), J, 3*(NL + 2) + 2, byrow = TRUE)
  # Time allocation
  ta <- FindTstar(tstar, g, hhopts)
  td <- cbind(matrix(8, J, 3*(NL + 2)), 0, 0)
  tb <- 16 - ta - td
  la <- To - ha - ta
  lb <- To - hb - tb - tcc*(gmat != gN & gmat > 0)
  lb[sm == 1, ncol(lb)-1] <- NA
  # Child quality
  optimal_cq <- ChildQ(q0, ta, tb, gmat, g0i, gf, gm, gmh, gmc, gsm, re, ed)
  initial_cq <- matrix(q0, J, ncol(optimal_cq))
  # Shares
  u          <- eq$u
  cons       <- eq$cons
  share_new  <- eq$probs
  s          <- SPEQ$s
  # Wages offer
  wa <- as.numeric(wi[, 1])
  wb <- as.numeric(wi[, 2])
  
  # _________________________________________________________________________________ ####
  # SIMULATED MOMENTS ####
  #' Computes the simulated moments that correspond to the vector of sample moments
  
  ## Column indicator
  m  <- seq(1, 3*(NL + 2), by = 3)
  mo <- 3*NL + 7
  
  ## Proportions
  # Married, College+
  i         <- which(sm == 0 & ed == 3)
  prop_1_cm <- sum(share_new[i, m])/sum(share_new[i, ])
  prop_2_cm <- sum(share_new[i, c(m + 1, mo)])/sum(share_new[i, ])
  prop_3_cm <- 1 - (prop_1_cm + prop_2_cm)
  # Married, <College
  i         <- which(sm == 0 & ed <= 2)
  prop_1_hm <- sum(share_new[i, m])/sum(share_new[i, ])
  prop_2_hm <- sum(share_new[i, c(m + 1, mo)])/sum(share_new[i, ])
  prop_3_hm <- 1 - (prop_1_hm + prop_2_hm)
  # Single, All
  i         <- which(sm == 1)
  prop_1_sm <- sum(share_new[i, m])/sum(share_new[i, ])
  prop_2_sm <- sum(share_new[i, c(m + 1, mo)])/sum(share_new[i, ])
  prop_3_sm <- 1 - (prop_1_sm + prop_2_sm)
  
  ## Relative care
  k <- which(gmat[1, ] == gR)
  # Married, College+
  i           <- which(sm == 0 & ed == 3)
  propR_12_cm <- sum(share_new[i, k[1:2]])/sum(share_new[i, c(m, m + 1, mo)])
  propR_3_cm  <- sum(share_new[i, k[3]])/sum(share_new[i, c(m + 2, mo + 1)])
  # Married, <College
  i           <- which(sm == 0 & ed <= 2)
  propR_12_hm <- sum(share_new[i, k[1:2]])/sum(share_new[i, c(m, m + 1, mo)])
  propR_3_hm  <- sum(share_new[i, k[3]])/sum(share_new[i, c(m + 2, mo + 1)])
  # Single, All
  i           <- which(sm == 1)
  propR_12_sm <- sum(share_new[i, k[1:2]])/sum(share_new[i, c(m, m + 1, mo)])
  propR_3_sm  <- sum(share_new[i, k[3]])/sum(share_new[i, c(m + 2, mo + 1)])
  
  ## Non-Relative care
  k <- which(gmat[1, ] == gN)
  # Married, College+
  i           <- which(sm == 0 & ed == 3)
  propN_1_cm  <- sum(share_new[i, k[1]])/sum(share_new[i, m])
  propN_2_cm  <- sum(share_new[i, k[2]])/sum(share_new[i, c(m + 1, mo)])
  propN_3_cm  <- sum(share_new[i, k[3]])/sum(share_new[i, c(m + 2, mo + 1)])
  # Married, <College
  i           <- which(sm == 0 & ed <= 2)
  propN_hm    <- sum(share_new[i, k])/sum(share_new[i, ])
  # Single, All
  i           <- which(sm == 1)
  propN_sm    <- sum(share_new[i, k])/sum(share_new[i, ])
  
  ## Center care
  k <- m[1:NL]
  # Married, College+
  i          <- which(sm == 0 & ed == 3)
  propC_1_cm <- sum(share_new[i, k])/sum(share_new[i, m])
  propC_2_cm <- sum(share_new[i, k + 1])/sum(share_new[i, c(m + 1, mo)])
  propC_3_cm <- sum(share_new[i, k + 2])/sum(share_new[i, c(m + 2, mo + 1)])
  # Married, <College
  i          <- which(sm == 0 & ed <= 2)
  propC_1_hm <- sum(share_new[i, k])/sum(share_new[i, m])
  propC_2_hm <- sum(share_new[i, k + 1])/sum(share_new[i, c(m + 1, mo)])
  propC_3_hm <- sum(share_new[i, k + 2])/sum(share_new[i, c(m + 2, mo + 1)])
  # Single, All
  i          <- which(sm == 1)
  propC_1_sm <- sum(share_new[i, k])/sum(share_new[i, m])
  propC_3_sm <- sum(share_new[i, k + 2])/sum(share_new[i, c(m + 2, mo + 1)])
  
  ## HQ vs. LQ usage %
  k  <- m[1:NL]
  k1 <- m[which(g == gH)]
  k2 <- which(gmat[1, ] == gH)
  # Married, College+
  i            <- which(sm == 0 & ed == 3)
  propHQ_12_cm <- sum(share_new[i, c(k1, k1 + 1)])/sum(share_new[i, c(k, k + 1)])
  propHQ_cm    <- sum(share_new[i, k2])/sum(share_new[i, 1:(3*NL)])
  # Married, <College
  i            <- which(sm == 0 & ed <= 2)
  propHQ_12_hm <- sum(share_new[i, c(k1, k1 + 1)])/sum(share_new[i, c(k, k + 1)])
  propHQ_hm    <- sum(share_new[i, k2])/sum(share_new[i, 1:(3*NL)])
  # Single, All
  i            <- which(sm == 1)
  propHQ_12_sm <- sum(share_new[i, c(k1, k1 + 1)])/sum(share_new[i, c(k, k + 1)])
  propHQ_sm    <- sum(share_new[i, k2])/sum(share_new[i, 1:(3*NL)])
  
  ## Center Price
  k  <- m[1:NL]
  sh <- share_new[, k] + share_new[, k + 1] + share_new[, k + 2]
  # Married, College+
  i        <- which(sm == 0 & ed == 3)
  meanp_cm <- weighted.mean(p[1:NL], colSums(sh[i, ]))
  sdp_cm   <- sqrt(weighted.mean((p[1:NL] - meanp_cm)^2, colSums(sh[i, ])))
  # Married, <College
  i        <- which(sm == 0 & ed <= 2)
  meanp_hm <- weighted.mean(p[1:NL], colSums(sh[i, ]))
  sdp_hm   <- sqrt(weighted.mean((p[1:NL] - meanp_hm)^2, colSums(sh[i, ])))
  # Single, All
  i        <- which(sm == 1)
  meanp_sm <- weighted.mean(p[1:NL], colSums(sh[i, ]))
  sdp_sm   <- sqrt(weighted.mean((p[1:NL] - meanp_sm)^2, colSums(sh[i, ])))
  # HQ
  k        <- which(g == gH)
  meanp_hq <- weighted.mean(p[k], colSums(as.matrix(sh[, k])))
  sdp_hq   <- sqrt(weighted.mean((p[k] - meanp_hq)^2, colSums(as.matrix(sh[, k]))))
  # LQ
  k        <- which(g == gL)
  meanp_lq <- weighted.mean(p[k], colSums(as.matrix(sh[, k])))
  sdp_lq   <- sqrt(weighted.mean((p[k] - meanp_lq)^2, colSums(as.matrix(sh[, k]))))
  
  ## Dad Wages
  # Married, College+
  i         <- which(sm == 0 & ed == 3)
  w         <- rowSums(share_new[i, m])
  wf_1_cm   <- weighted.mean(wa[i], w)
  swf_1_cm  <- sqrt(weighted.mean((wa[i] - wf_1_cm)^2, w))
  w         <- rowSums(share_new[i, c(m + 1, mo)])
  wf_2_cm   <- weighted.mean(wa[i], w)
  swf_2_cm  <- sqrt(weighted.mean((wa[i] - wf_2_cm)^2, w))
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wf_12_cm  <- weighted.mean(wa[i], w)
  swf_12_cm <- sqrt(weighted.mean((wa[i] - wf_12_cm)^2, w))
  w         <- rowSums(share_new[i, c(m + 2, mo + 1)])
  wf_3_cm   <- weighted.mean(wa[i], w)
  swf_3_cm  <- sqrt(weighted.mean((wa[i] - wf_3_cm)^2, w))
  # Married, <College
  i         <- which(sm == 0 & ed <= 2)
  w         <- rowSums(share_new[i, m])
  wf_1_hm   <- weighted.mean(wa[i], w)
  swf_1_hm  <- sqrt(weighted.mean((wa[i] - wf_1_hm)^2, w))
  w         <- rowSums(share_new[i, c(m + 1, mo)])
  wf_2_hm   <- weighted.mean(wa[i], w)
  swf_2_hm  <- sqrt(weighted.mean((wa[i] - wf_2_hm)^2, w))
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wf_12_hm  <- weighted.mean(wa[i], w)
  swf_12_hm <- sqrt(weighted.mean((wa[i] - wf_12_hm)^2, w))
  w         <- rowSums(share_new[i, ])
  wf_hm     <- weighted.mean(wa[i], w)
  swf_hm    <- sqrt(weighted.mean((wa[i] - wf_hm)^2, w))
  
  ## Mom Wages
  # Married, College+
  i         <- which(sm == 0 & ed == 3)
  w         <- rowSums(share_new[i, m])
  wm_1_cm   <- weighted.mean(wb[i], w)
  swm_1_cm  <- sqrt(weighted.mean((wb[i] - wm_1_cm)^2, w))
  w         <- rowSums(share_new[i, c(m + 1, mo)])
  wm_2_cm   <- weighted.mean(wb[i], w)
  swm_2_cm  <- sqrt(weighted.mean((wb[i] - wm_2_cm)^2, w))
  # Married, <College
  i         <- which(sm == 0 & ed <= 2)
  w         <- rowSums(share_new[i, m])
  wm_1_hm   <- weighted.mean(wb[i], w)
  swm_1_hm  <- sqrt(weighted.mean((wb[i] - wm_1_hm)^2, w))
  w         <- rowSums(share_new[i, c(m + 1, mo)])
  wm_2_hm   <- weighted.mean(wb[i], w)
  swm_2_hm  <- sqrt(weighted.mean((wb[i] - wm_2_hm)^2, w))
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_hm  <- weighted.mean(wb[i], w)
  swm_12_hm <- sqrt(weighted.mean((wb[i] - wm_12_hm)^2, w))
  # Single, All
  i         <- which(sm == 1)
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_sm  <- weighted.mean(wb[i], w)
  swm_12_sm <- sqrt(weighted.mean((wb[i] - wm_12_sm)^2, w))
  # All, College+
  i         <- which(ed == 3)
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_c   <- weighted.mean(wb[i], w)
  swm_12_c  <- sqrt(weighted.mean((wb[i] - wm_12_c)^2, w))
  # All, HS+
  i         <- which(ed == 2)
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_hp  <- weighted.mean(wb[i], w)
  swm_12_hp <- sqrt(weighted.mean((wb[i] - wm_12_hp)^2, w))
  # All, <HS
  i         <- which(ed == 1)
  w         <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_dp  <- weighted.mean(wb[i], w)
  swm_12_dp <- sqrt(weighted.mean((wb[i] - wm_12_dp)^2, w))
  # Married, HS+
  i          <- which(ed == 2 & sm == 0)
  w          <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_hpm  <- weighted.mean(wb[i], w)
  swm_12_hpm <- sqrt(weighted.mean((wb[i] - wm_12_hpm)^2, w))
  # Single, College+
  i          <- which(ed == 3 & sm == 1)
  w          <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_cs   <- weighted.mean(wb[i], w)
  swm_12_cs  <- sqrt(weighted.mean((wb[i] - wm_12_cs)^2, w))
  # Single, HS+
  i          <- which(ed == 2 & sm == 1)
  w          <- rowSums(share_new[i, c(m, m + 1, mo)])
  wm_12_hps  <- weighted.mean(wb[i], w)
  swm_12_hps <- sqrt(weighted.mean((wb[i] - wm_12_hps)^2, w))
  
  ## Corr Wages
  # Married, College+
  i         <- which(sm == 0 & ed == 3)
  cov_wab   <- (wa[i] - wf_1_cm) * (wb[i] - wm_1_cm)
  cov_wab   <- weighted.mean(cov_wab, rowSums(share_new[i, m]))
  corw_1_cm <- cov_wab / (swf_1_cm * swm_1_cm)
  cov_wab   <- (wa[i] - wf_2_cm) * (wb[i] - wm_2_cm)
  cov_wab   <- weighted.mean(cov_wab, rowSums(share_new[i, c(m + 1, mo)]))
  corw_2_cm <- cov_wab / (swf_2_cm * swm_2_cm)
  # Married, <College
  i         <- which(sm == 0 & ed <= 2)
  cov_wab   <- (wa[i] - wf_1_hm) * (wb[i] - wm_1_hm)
  cov_wab   <- weighted.mean(cov_wab, rowSums(share_new[i, m]))
  corw_1_hm <- cov_wab / (swf_1_hm * swm_1_hm)
  cov_wab   <- (wa[i] - wf_2_hm) * (wb[i] - wm_2_hm)
  cov_wab   <- weighted.mean(cov_wab, rowSums(share_new[i, c(m + 1, mo)]))
  corw_2_hm <- cov_wab / (swf_2_hm * swm_2_hm)
  
  ## Mental Scores
  # Married, College+
  i         <- which(sm == 0 & ed == 3)
  meanq_cm  <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # Married, <College
  i         <- which(sm == 0 & ed <= 2)
  meanq_hm  <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # Married, All
  i          <- which(sm == 0)
  meanq_3_m  <- weighted.mean(optimal_cq[i, c(m + 2, mo + 1)], 
                              share_new[i, c(m + 2, mo + 1)])
  # Single, All
  i          <- which(sm == 1)
  meanq_3_sm <- weighted.mean(optimal_cq[i, c(m + 2, mo + 1)], 
                              share_new[i, c(m + 2, mo + 1)])
  meanq_sm   <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # All
  meanq_12  <- weighted.mean(optimal_cq[, c(m, m + 1, mo)], 
                             share_new[, c(m, m + 1, mo)])
  meanq_3   <- weighted.mean(optimal_cq[, c(m + 2, mo + 1)], 
                             share_new[, c(m + 2, mo + 1)])
  # All, College+
  i         <- which(ed == 3)
  meanq_c   <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # All, HS+
  i         <- which(ed == 2)
  meanq_hp  <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # All, <HS
  i         <- which(ed == 1)
  meanq_dp  <- weighted.mean(optimal_cq[i, ], share_new[i, ])
  # Married, HS+
  i           <- which(ed == 2 & sm == 0)
  meanq_3_hpm <- weighted.mean(optimal_cq[i, c(m + 2, mo + 1)], 
                               share_new[i, c(m + 2, mo + 1)])
  # Single, HS+
  i           <- which(ed == 2 & sm == 1)
  meanq_3_hps <- weighted.mean(optimal_cq[i, c(m + 2, mo + 1)], 
                               share_new[i, c(m + 2, mo + 1)])
  # All, type of care
  k        <- which(gmat[1, ] == gH)
  meanq_hq <- weighted.mean(optimal_cq[, k], share_new[, k])
  k        <- which(gmat[1, ] == gL)
  meanq_lq <- weighted.mean(optimal_cq[, k], share_new[, k])
  k        <- which(gmat[1, ] == gR)
  meanq_R  <- weighted.mean(optimal_cq[, k], share_new[, k])
  k        <- which(gmat[1, ] == gN)
  meanq_N  <- weighted.mean(optimal_cq[, k], share_new[, k])
  
  # _________________________________________________________________________________ ####
  # FINAL MOMENTS ####
  #' Final outcomes of the function: 
  #' simulated vector of moments, implied unobserved heterogeneity sources
  #' and equilibrium outcome variables
  
  ## Final moment vector
  smoms_fin <- c(prop_1_cm*100, propN_1_cm*100, propC_1_cm*100, 
                 wm_1_cm, swm_1_cm, corw_1_cm*100, 
                 prop_2_cm*100, propN_2_cm*100, propC_2_cm*100, 
                 wm_2_cm, swm_2_cm, corw_2_cm*100,
                 prop_3_cm*100, propR_3_cm*100, propN_3_cm*100, propC_3_cm*100, 
                 wf_3_cm, swf_3_cm, 
                 meanp_cm, sdp_cm, meanq_cm, propHQ_cm*100,
                 propR_12_cm*100, wf_12_cm, swf_12_cm, propHQ_12_cm*100,
                 prop_1_hm*100, propC_1_hm*100, corw_1_hm*100,
                 prop_2_hm*100, propC_2_hm*100, corw_2_hm*100,
                 prop_3_hm*100, propR_3_hm*100, propC_3_hm*100,
                 propN_hm*100, meanp_hm, sdp_hm, 
                 wf_hm, swf_hm, wm_12_hm, swm_12_hm,
                 meanq_hm, propHQ_hm*100,
                 propR_12_hm*100, wf_12_hm, swf_12_hm, propHQ_12_hm*100,
                 prop_1_sm*100, propC_1_sm*100, 
                 prop_3_sm*100, propR_3_sm*100, propC_3_sm*100,
                 propN_sm*100, meanp_sm, sdp_sm, 
                 wm_12_sm, swm_12_sm, meanq_sm, propHQ_sm*100,
                 propR_12_sm*100, propHQ_12_sm*100,
                 meanp_hq, sdp_hq, meanp_lq, sdp_lq,
                 meanq_hq, meanq_lq, meanq_R, meanq_N,
                 meanq_12, meanq_3, 
                 wm_12_c, swm_12_c, wm_12_hp, swm_12_hp, wm_12_dp, swm_12_dp,
                 wm_12_hp, swm_12_hp,
                 wm_12_cs, swm_12_cs, wm_12_hps, swm_12_hps,
                 meanq_c, meanq_hp, meanq_dp, 
                 meanq_3_hpm, meanq_3_hps, meanq_3_m, meanq_3_sm) 
  
  smoms_fin <- smoms_fin[-nomoms]
  
  return(list(smoms_fin = smoms_fin, SPEQ = SPEQ, u = u, cons = cons, share_new = share_new, 
              ta = ta, tb = tb, la = la, lb = lb, ha = ha, hb = hb, gmat = gmat,
              oq = optimal_cq, alphas = ai, pars = pars, Cg = Cg, wi = wi))
  
}
